import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import fsolve
import pandas as pd
#

# 设置matplotlib中文字体支持
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False  


class RamseyModel:
    """
    拉姆齐-卡斯-库普曼斯模型实现
    """
    
    def __init__(self, alpha=0.3, beta=0.95, delta=0.1, rho=0.02, sigma=1.0, A=1.0):
        """
        初始化模型参数
        """
        self.alpha = alpha      # 资本产出弹性
        self.beta = beta        # 折现因子
        self.delta = delta      # 折旧率
        self.rho = rho          # 时间偏好率
        self.sigma = sigma      # 相对风险厌恶系数
        self.A = A              # 技术水平
        
        # 计算稳态值
        self.k_star, self.c_star = self.calculate_steady_state()
    
    def production_function(self, k):
        """柯布-道格拉斯生产函数"""
        return self.A * k**self.alpha
    
    def marginal_product_capital(self, k):
        """资本的边际产出"""
        return self.alpha * self.A * k**(self.alpha - 1)
    
    def utility(self, c):
        """CRRA效用函数"""
        if self.sigma == 1:
            return np.log(c)
        else:
            return (c**(1 - self.sigma) - 1) / (1 - self.sigma)
    
    def marginal_utility(self, c):
        """边际效用"""
        return c**(-self.sigma)
    
    def calculate_steady_state(self):
        """计算稳态资本和消费"""
        # 稳态条件: r = ρ + δ
        k_star = ((self.alpha * self.A) / (self.rho + self.delta))**(1 / (1 - self.alpha))
        c_star = self.production_function(k_star) - self.delta * k_star
        return k_star, c_star
    
    def consumption_growth_rate(self, k, c):
        """消费增长率 (Keynes-Ramsey规则)"""
        r = self.marginal_product_capital(k) - self.delta
        return (r - self.rho) / self.sigma
    
    def capital_accumulation(self, k, c):
        """资本积累方程"""
        return self.production_function(k) - c - self.delta * k
    
    def phase_diagram_k(self, k):
        """相图中dk/dt=0的曲线"""
        return self.production_function(k) - self.delta * k
    
    def phase_diagram_c(self, k):
        """相图中dc/dt=0的曲线"""
        return self.production_function(k) - self.delta * k  # 简化版本
    
    def solve_trajectory(self, k0, c0_guess, T=100, dt=0.1):
        """
        求解从给定初始资本出发的最优路径
        使用打靶法(shooting method)找到最优初始消费
        """
        def objective(c0):
            """目标函数: 最终资本与稳态资本的差距"""
            k_path, c_path = self.simulate_forward(k0, c0, T, dt)
            return k_path[-1] - self.k_star
        
        # 找到最优初始消费
        c0_optimal = fsolve(objective, c0_guess)[0]
        
        # 用最优初始消费模拟路径
        k_path, c_path = self.simulate_forward(k0, c0_optimal, T, dt)
        
        return k_path, c_path, c0_optimal
    
    def simulate_forward(self, k0, c0, T, dt):
        """向前模拟经济路径"""
        steps = int(T / dt)
        k_path = np.zeros(steps)
        c_path = np.zeros(steps)
        
        k_path[0] = k0
        c_path[0] = c0
        
        for t in range(1, steps):
            # 欧拉方法更新
            k_dot = self.capital_accumulation(k_path[t-1], c_path[t-1])
            c_dot = c_path[t-1] * self.consumption_growth_rate(k_path[t-1], c_path[t-1])
            
            k_path[t] = k_path[t-1] + k_dot * dt
            c_path[t] = c_path[t-1] + c_dot * dt
            
            # 确保消费非负
            c_path[t] = max(c_path[t], 0.01)
        
        return k_path, c_path
    
    def plot_phase_diagram(self, k_min=0.1, k_max=2.0):
        """绘制相图"""
        k_values = np.linspace(k_min, k_max, 100)
        
        # dk/dt = 0 的曲线
        c_dk_zero = self.phase_diagram_k(k_values)
        
        # dc/dt = 0 的条件 (r = ρ)
        k_dc_zero = ((self.alpha * self.A) / (self.rho + self.delta))**(1/(1-self.alpha))
        c_dc_zero = self.phase_diagram_k(k_dc_zero)
        
        plt.figure(figsize=(12, 8))
        
        # 相图背景
        k_grid, c_grid = np.meshgrid(np.linspace(k_min, k_max, 20), 
                                    np.linspace(0.1, 1.5, 20))
        
        k_dot = self.capital_accumulation(k_grid, c_grid)
        c_dot = c_grid * self.consumption_growth_rate(k_grid, c_grid)
        
        plt.streamplot(k_grid, c_grid, k_dot, c_dot, density=1.5, 
                      color='lightgray', linewidth=0.5)
        
        # 绘制零变化线
        plt.plot(k_values, c_dk_zero, 'b-', linewidth=2, label='$\dot{k}=0$')
        plt.axvline(x=k_dc_zero, color='red', linestyle='--', 
                   label='$\dot{c}=0$', linewidth=2)
        
        # 标记稳态
        plt.plot(self.k_star, self.c_star, 'ro', markersize=10, 
                label=f'稳态 ($k^*=${self.k_star:.2f}, $c^*=${self.c_star:.2f})')
        
        # 示例轨迹
        initial_conditions = [
            (self.k_star * 0.5, '低于稳态'),
            (self.k_star * 1.5, '高于稳态')
        ]
        
        colors = ['green', 'purple']
        for i, (k0, label) in enumerate(initial_conditions):
            k_path, c_path, c0_opt = self.solve_trajectory(k0, self.c_star * 0.8)
            plt.plot(k_path, c_path, color=colors[i], linewidth=2, 
                    label=f'最优路径 ({label})')
            plt.plot(k0, c0_opt, 'o', color=colors[i], markersize=6)
        
        plt.xlabel('资本 k', fontsize=12)
        plt.ylabel('消费 c', fontsize=12)
        plt.title('拉姆齐模型相图', fontsize=14)
        plt.legend()
        plt.grid(True, alpha=0.3)
        # plt.xlim(k_min, k_max)
        plt.xlim(k_min, 4)
        plt.ylim(0, 1.5)
        plt.tight_layout()
        plt.show()
    
    def plot_time_paths(self, k0):
        """绘制时间路径"""
        k_path, c_path, c0_opt = self.solve_trajectory(k0, self.c_star * 0.8)
        time = np.arange(len(k_path)) * 0.1
        
        # 自动计算合适的坐标轴范围
        x_margin = 5  # 时间轴额外范围
        k_margin = 0.2 * (max(k_path) - min(k_path))  # 资本轴20%边距
        c_margin = 0.2 * (max(c_path) - min(c_path))  # 消费轴20%边距
        
        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 10))
        
        # 资本路径
        ax1.plot(time, k_path, 'b-', linewidth=2, label='资本 k')
        ax1.axhline(y=self.k_star, color='r', linestyle='--', label='稳态资本')
        ax1.set_ylabel('资本', fontsize=12)
        ax1.set_title('资本积累路径', fontsize=14)
        ax1.legend()
        ax1.grid(True, alpha=0.3)
        ax1.set_xlim(0, max(time) + x_margin)
        ax1.set_ylim(min(k_path) - k_margin, max(k_path) + k_margin)
        
        # 消费路径
        ax2.plot(time, c_path, 'g-', linewidth=2, label='消费 c')
        ax2.axhline(y=self.c_star, color='r', linestyle='--', label='稳态消费')
        ax2.set_xlabel('时间', fontsize=12)
        ax2.set_ylabel('消费', fontsize=12)
        ax2.set_title('消费演化路径', fontsize=14)
        ax2.legend()
        ax2.grid(True, alpha=0.3)
        ax2.set_xlim(0, max(time) + x_margin)
        ax2.set_ylim(min(c_path) - c_margin, max(c_path) + c_margin)
        
        plt.tight_layout()
        plt.show()
        
        return k_path, c_path

# 示例使用
if __name__ == "__main__":
    # 创建模型实例（使用典型参数）
    model = RamseyModel(alpha=0.3, beta=0.95, delta=0.1, rho=0.02, sigma=1.0)
    
    print("模型参数:")
    print(f"资本产出弹性 α = {model.alpha}")
    print(f"折旧率 δ = {model.delta}")
    print(f"时间偏好率 ρ = {model.rho}")
    print(f"相对风险厌恶系数 σ = {model.sigma}")
    print(f"稳态资本 k* = {model.k_star:.4f}")
    print(f"稳态消费 c* = {model.c_star:.4f}")
    
    # 绘制相图
    print("\n生成相图...")
    model.plot_phase_diagram()
    
    # 绘制时间路径示例
    print("生成时间路径...")
    k0_example = model.k_star * 0.6  # 初始资本低于稳态
    k_path, c_path = model.plot_time_paths(k0_example)
    
    # 敏感性分析：不同参数的影响
    print("\n敏感性分析:")
    param_variations = {
        '高时间偏好': {'rho': 0.05},
        '低时间偏好': {'rho': 0.01},
        '高风险厌恶': {'sigma': 2.0},
        '低风险厌恶': {'sigma': 0.5}
    }
    
    results = []
    for name, params in param_variations.items():
        temp_model = RamseyModel(**params)
        results.append({
            'Scenario': name,
            'k_star': temp_model.k_star,
            'c_star': temp_model.c_star,
            'Params': params
        })
    
    results_df = pd.DataFrame(results)
    print(results_df[['Scenario', 'k_star', 'c_star']])